package cn.trigram.mgb;

import cn.trigram.mgb.ext.AbstractContextCustomize;
import cn.trigram.mgb.ext.AbstractContextCustomize.DefaultContextCustomize;
import cn.trigram.mgb.ext.AbstractTableGenerateCustomize;
import cn.trigram.mgb.ext.AbstractTableGenerateCustomize.DefaultTableGenerateCustomize;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import org.mybatis.generator.config.Context;

public class ThreadLocalUtil {

  private static final ThreadLocal<Map<String, AbstractTableGenerateCustomize>> TABLE_GENERATE_CUSTOMIZE_THREAD_LOCAL = ThreadLocal.withInitial(
      HashMap::new);

  private static final ThreadLocal<Map<String, AbstractContextCustomize>> CONTEXT_CUSTOMIZE_THREAD_LOCAL = ThreadLocal.withInitial(
      HashMap::new);

  public static void addTableGenerateCustomize(Context context, AbstractTableGenerateCustomize tableGenerateCustomize) {

    TABLE_GENERATE_CUSTOMIZE_THREAD_LOCAL.get().put(context.getId(), tableGenerateCustomize);
  }

  public static AbstractTableGenerateCustomize getTableGenerateCustomize(Context context) {

    AbstractTableGenerateCustomize tableGenerateCustomize = TABLE_GENERATE_CUSTOMIZE_THREAD_LOCAL.get()
        .get(context.getId());

    return Optional.ofNullable(tableGenerateCustomize).orElseGet(DefaultTableGenerateCustomize::new);
  }

  public static void addContextCustomize(Context context, AbstractContextCustomize contextCustomize) {

    CONTEXT_CUSTOMIZE_THREAD_LOCAL.get().put(context.getId(), contextCustomize);
  }

  public static AbstractContextCustomize getContextCustomize(Context context) {

    AbstractContextCustomize contextCustomize = CONTEXT_CUSTOMIZE_THREAD_LOCAL.get()
        .get(context.getId());

    return Optional.ofNullable(contextCustomize).orElseGet(DefaultContextCustomize::new);
  }

}
